package com.tsx.china.mobile.ai.controller;

import com.tsx.china.mobile.ai.service.ChatService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.util.List;

@RestController
@Slf4j
@RequiredArgsConstructor
@RequestMapping("/china-mobile")
public class ChatController {
    private final ChatService chatService;
    private final VectorStore vectorStore;

    @PostMapping(value = "/chat/{sessionId}", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> chat(@PathVariable String sessionId, @RequestBody String question) {
        return chatService.chat(sessionId, question);
    }

    @PostMapping("/search")
    public List<Document> search(@RequestParam("query") String query) {
        return this.vectorStore.similaritySearch(SearchRequest.builder()
                .query(query) // 设置查询条件
                .topK(3) // 设置最多返回的文档数量
                .build());
    }
}
